#include "bp.h"

#include <iomanip>
#include <float.h>

namespace BP {

double
log_sum(double a, double b)
{
	if (a>b) return a + log(1+exp(b-a));
	else	 return b + log(1+exp(a-b));
}

void 
calculate_dot_products(EventPtr e, boost::shared_array<double> weights, 
	Array3d &nodes, Array4d &cliques, bool generic_cliques)
{ calculate_dot_products(e, weights.get(), nodes, cliques, generic_cliques); }

void
calculate_dot_products(EventPtr e, 
//	boost::shared_array<double> weights, 
	double *weights, 
	Array3d &nodes, Array4d &cliques, 
	bool generic_cliques)
{
	int event_length = e->length;
	SparseArrayPtr sa;
	for (int n=0; n < event_length; ++n)
	{
		for (int o=0; (n-o)>=0 && o<(int)e->order; ++o)
		{
			int start = n-o;
			NodePtr node = e->nodes[o][start];
			int num_labels = node->features.size();
			for (int b=0; b<num_labels; ++b)
			{
// 				std::cerr << n << " " << o << " " << b << std::endl;
				double dot=0.0;
				sa = node->features[b];
				for(int j=0; j < sa->active(); ++j)
					dot += weights[sa->_index[j]]*sa->_values[j];
				nodes[n][o][b] = dot;
// 				std::cerr << "\t" << dot << std::endl;

				if (start!=0)
				{
					// the generic_cliques flag means we only calculate features for one clique, and
					// assume all others are the same
					if (!generic_cliques || start==1)
					{
						BaseCliquePtr clique = e->cliques[o][start-1];
						int prev_num_labels = clique->labels().first;
						assert(clique->labels().second==num_labels);
						for (int a=0; a<prev_num_labels; ++a)
						{
							dot=0.0;
							sa = clique->labelling(a,b);
							for(int j=0; j < sa->active(); ++j)
							{
								dot += weights[sa->_index[j]]*sa->_values[j];
//								std::cerr << sa->_index[j] << a << " " << b << std::endl;
							}
							cliques[n-1][o][a][b] = dot;
	//						cliques[n-1][o][a][b] = 0;
//							std::cerr << a << " " << b << " " << dot << std::endl;
//							std::cerr << *sa << "=" << dot << std::endl;

						}
					}
				}
			}
		}
	}
}

LabellingPtr
viterbi(EventPtr v, Array3d &nodes, Array4d &cliques, bool print, bool generic_cliques)
{
//	print = true;
	// assume every labelled clique is the same size
	size_t global_width = nodes.shape()[2];;
	size_t num_nodes = v->length;

	// data structures for max products and back pointers
	Array2d V(boost::extents[num_nodes][global_width]);
	boost::multi_array<std::pair<int,size_t>, 2> B(boost::extents[num_nodes][global_width]);

	/////////////////////////////////////////////////////////////////
	// main max product(forward) loop
	if (print) std::cout << "Viterbi lattice:" << std::endl;
	for (int i=0; i < (int)num_nodes; ++i)
	{
		size_t width = v->nodes[0][i]->features.size(); // not valid for semi-CRF
		if (v->order>1) assert(global_width==width);
		for (size_t b=0; b < width; ++b)
			V[i][b] = -FLT_MAX;

		double potential=0.0;
		for (int o=0; (i-o)>=0 && o<(int)v->order; ++o)
		{
			int start = i-o;
			for (size_t b=0; b < width; ++b)
			{
				// node_contribution only
				if (start == 0)
				{
					potential = nodes[i][o][b];

					// MAX PRODUCT
					if (potential > V[i][b])
					{
						V[i][b] = potential;
						B[i][b] = std::pair<int,size_t>(-1,0);
					}
				}	
				// clique and nodes contribution
				else
				{ 
					// not valid for semi-CRF
					size_t prev_width = v->nodes[0][start-1]->features.size();

					int clique_index=i-1;
					if (generic_cliques) clique_index=0; // just reuse the same clique features
					assert(clique_index>=0);

					for (size_t a=0; a < prev_width; ++a)
					{
						potential = V[start-1][a]+cliques[clique_index][o][a][b]
									+nodes[i][o][b];

						// MAX PRODUCT
						if (potential > V[i][b])
						{
							V[i][b] = potential;
							B[i][b] = std::pair<int,size_t>(start-1,a);
						}
					}
				}
			}
		}
		if (print)
		{
			for (size_t j=0; j < width; ++j)
				std::cout << " " << std::setprecision(3) << std::setw(6) 
					<< V[i][j];
			std::cout << std::endl;
			int clique_index=i-1;
			if (generic_cliques) clique_index=0; assert(clique_index>=0);
			for (size_t j=0; j < width; ++j)
//				std::cout << " " << std::setprecision(3) << std::setw(6) 
//					<< B[i][j].first << "-" << B[i][j].second;
				if (i>0)
				std::cout << " " << std::setprecision(3) << std::setw(6) 
					<< cliques[clique_index][0][B[i][j].second][j] << "-" << B[i][j].second;
			std::cout << std::endl;
		}
	}
	if (print) std::cout << std::endl;

	/////////////////////////////////////////////////////////////////
	// resolve max labelling 
	LabellingPtr result(new Labelling);
	double best_p = -FLT_MAX;
	size_t best_k = 0;
	size_t end_width = v->nodes[0][num_nodes-1]->features.size(); // not valid for semi-CRF
	for (size_t k=0; k < end_width; ++k)
		if (V[num_nodes-1][k] > best_p)
		{
			best_k = k;
			best_p = V[num_nodes-1][k];
		}
	result->push_front(std::pair<int,size_t>(num_nodes-1,best_k));

	// trace the back pointers
	std::pair<int,size_t> back_ptr = B[num_nodes-1][best_k];
	while(back_ptr.first >= 0)
	{
		result->push_front(back_ptr);
		back_ptr = B[back_ptr.first][back_ptr.second];
	}

	if (print)
	{
		for (size_t j=0; j < result->size(); ++j)
			  std::cout << " " << std::setprecision(2) << std::setw(8) << j;
		std::cout << std::endl;
		for (Labelling::iterator it=result->begin(); it != result->end(); ++it)
			  std::cout << " " << std::setprecision(2) << std::setw(8) 
				<< it->first << "-" << it->second;
		std::cout << std::endl;
		std::cout << std::endl;
		std::cout << std::endl;
	}
	return result;
}

LabellingPtr
max_cost_viterbi(EventPtr v, Array3d &nodes, Array4d &cliques, bool print,
	double (*cost_func)(int,double,int,double))
{
//	print = true;
	// assume every labelled clique is the same size
	size_t global_width = nodes.shape()[2];;
	size_t num_nodes = v->length;

	// data structures for max products and back pointers
	Array2d V(boost::extents[num_nodes][global_width]);
	boost::multi_array<std::pair<int,size_t>, 2>
		B(boost::extents[num_nodes][global_width]);

	/////////////////////////////////////////////////////////////////
	// main max product(forward) loop
	if (print) std::cout << "Viterbi lattice:" << std::endl;
	for (int i=0; i < (int)num_nodes; ++i)
	{
		size_t width = v->nodes[0][i]->features.size(); // not valid for semi-CRF
		if (v->order>1) assert(global_width==width);
		for (size_t b=0; b < width; ++b)
			V[i][b] = -FLT_MAX;

		double potential=0.0;
		for (int o=0; (i-o)>=0 && o<(int)v->order; ++o)
		{
			int start = i-o;
			size_t gold=v->nodes[o][start]->gold;
			for (size_t b=0; b < width; ++b)
			{
				// node_contribution only
				if (start == 0)
				{
					potential = cost_func(b,nodes[i][o][b],gold,nodes[i][o][gold]);

					// MAX PRODUCT
					if (potential > V[i][b])
					{
						V[i][b] = potential;
						B[i][b] = std::pair<int,size_t>(-1,0);
					}
				}
				// clique and nodes contribution
				else
				{ 
					// not valid for semi-CRF
					size_t prev_width = v->nodes[0][start-1]->features.size();
					for (size_t a=0; a < prev_width; ++a)
					{
						double cost = cost_func(b,cliques[i-1][o][a][b]+nodes[i][o][b],
								gold,cliques[i-1][o][a][gold]+nodes[i][o][gold]);
						potential = V[start-1][a]+cost;

						// MAX PRODUCT
						if (potential > V[i][b])
						{
							V[i][b] = potential;
							B[i][b] = std::pair<int,size_t>(start-1,a);
						}
					}
				}
			}
		}
		if (print)
		{
			for (size_t j=0; j < width; ++j)
				std::cout << std::setprecision(3) << std::setw(9) << V[i][j];
			std::cout << std::endl;
			for (size_t j=0; j < width; ++j)
//				std::cout << " " << std::setprecision(3) << std::setw(6) 
//					<< B[i][j].first << "-" << B[i][j].second;
				if (i>0)
					std::cout << std::setprecision(3) << std::setw(7) 
						<< cliques[i-1][0][B[i][j].second][j] << "-" << B[i][j].second;
			std::cout << std::endl;
		}
	}
	if (print) std::cout << std::endl;

	/////////////////////////////////////////////////////////////////
	// resolve max labelling 
	LabellingPtr result(new Labelling);
	double best_p = -FLT_MAX;
	size_t best_k = 0;
	size_t end_width = v->nodes[0][num_nodes-1]->features.size(); // not valid for semi-CRF
	for (size_t k=0; k < end_width; ++k)
		if (V[num_nodes-1][k] > best_p)
		{
			best_k = k;
			best_p = V[num_nodes-1][k];
		}
	result->push_front(std::pair<int,size_t>(num_nodes-1,best_k));

	// trace the back pointers
	std::pair<int,size_t> back_ptr = B[num_nodes-1][best_k];
	while(back_ptr.first >= 0)
	{
		result->push_front(back_ptr);
		back_ptr = B[back_ptr.first][back_ptr.second];
	}

	if (print)
	{
		for (size_t j=0; j < result->size(); ++j)
			  std::cout << std::setprecision(2) << std::setw(8) << j;
		std::cout << std::endl;
		for (Labelling::iterator it=result->begin(); it != result->end(); ++it)
			  std::cout << std::setprecision(2) << std::setw(6) 
				<< it->first << "-" << it->second;
		std::cout << std::endl;
		std::cout << std::endl;
		std::cout << std::endl;
	}
	return result;
}
void
forward_backward(EventPtr v, Array3d &nodes, Array4d &cliques,
		Array2d &alpha, Array2d &beta, bool print)
{
	// assume every labelled clique is the same size
	size_t global_width = nodes.shape()[2];
	size_t num_nodes = v->length;

	/////////////////////////////////////////////////////////////////
	// main sum product(forward) loop
	if (print) std::cout << "Forward pass:" << std::endl;
	for (int i=0; i < (int)num_nodes; ++i)
	{
		size_t width = v->nodes[0][i]->features.size(); // not valid for semi-CRF
		if (v->order>1) assert(global_width==width);
		for (size_t b=0; b < width; ++b)
			alpha[i][b] = -FLT_MAX;

		for (int o=0; (i-o)>=0 && o<(int)v->order; ++o)
		{
			int start = i-o;
			for (size_t b=0; b < width; ++b)
			{
				if (start==0) // node_contribution only
					alpha[i][b] = log_sum(alpha[i][b], nodes[i][o][b]);
				else // clique and nodes contribution
				{
					size_t prev_width = v->nodes[0][start-1]->features.size();
					for (size_t a=0; a < prev_width; ++a)
						alpha[i][b] = log_sum(alpha[i][b],
							alpha[start-1][a]+cliques[i-1][o][a][b]+nodes[i][o][b]);
				}
			}
		}
		if (print)
		{
			for (size_t j=0; j < width; ++j)
				std::cout << " " << std::setprecision(3) << std::setw(6) 
					<< alpha[i][j];
			std::cout << std::endl;
		}
	}
	if (print) std::cout << std::endl;

	// main sum product(backward) loop
	// initialisation
	size_t width = v->nodes[0][num_nodes-1]->features.size(); // not valid for semi-CRF
	for (size_t i=0; i < width; ++i)
		beta[num_nodes-1][i] = 0.0; 

	if (print) std::cout << "Backward pass:" << std::endl;
	if (print)
	{
		for (size_t j=0; j < width; ++j)
			std::cout << " " << std::setprecision(2) << std::setw(6) 
				<< beta[num_nodes-1][j];
		std::cout << std::endl;
	}
	for (int i=num_nodes-2; i >= 0; --i)
	{
		width = v->nodes[0][i]->features.size(); // not valid for semi-CRF
		size_t next_width = v->nodes[0][i+1]->features.size(); // not valid for semi-CRF
		for (size_t a=0; a < width; a++)
		{
			beta[i][a] = -FLT_MAX;
			for (int o=0; (i+o) < (int)num_nodes-1 && o<(int)v->order; ++o)
			{
				int end = i+o+1;
				for (size_t b=0; b < next_width; b++)
				{
					beta[i][a] = log_sum(beta[i][a],
						beta[end][b]+cliques[end-1][o][a][b]+nodes[end][o][b]);
				}
			}
		}
		if (print)
		{
			for (size_t j=0; j < width; ++j)
					std::cout << " " << std::setprecision(3) << std::setw(6) << beta[i][j];
			std::cout << std::endl;
		}
	}
	if (print) std::cout << std::endl;

	return;
}

//std::pair<LatticePtr,boost::shared_ptr<std::vector<int> > >
//prob_lattice(EventPtr e, boost::shared_array<double> weights)
//{
//	  unsigned int num_cliques = e->cliques.size();
//
//	  std::vector<double> dot_products; 
//	  LatticePtr alpha, beta;
//	  calculate_dot_products(e, dot_products, weights);
//	  forward_backward(e, dot_products, alpha, beta);
//
//	  // create data structures
//	  boost::shared_ptr<std::vector<int> > 
//		  labels_ptr(new std::vector<int>(num_cliques+1));
//	  LatticePtr probs(new LatticeColumnPtr [num_cliques+1]);
//
//	  std::vector<int> &labels = *labels_ptr;
//
//	  e->cliques[0]->dimensions(labels[0], labels[1]);
//	  for (unsigned int i=1; i < num_cliques+1; ++i)
//		  e->cliques[i-1]->dimensions(labels[i-1], labels[i]);
//
//	  // calculate the probability of each label
//	  for (unsigned int i=0; i < num_cliques+1; ++i)
//	  {
//		  double Z=-FLT_MAX;
//		  for (int j=0; j < labels[i]; ++j)
//			  Z = log_sum(Z,alpha[i][j] + beta[i][j]);	
//
//		  probs[i] = LatticeColumnPtr(new double [labels[i]]);
//		  for (int j=0; j < labels[i]; ++j)
//			  probs[i][j] = exp(alpha[i][j]+beta[i][j] - Z);  
//	  }
//
//	  return std::pair<LatticePtr,boost::shared_ptr<std::vector<int> > >
//		  (LatticePtr(probs), labels_ptr);
//}

} // namespace BP
